#!/usr/bin/env python
# -*- coding: utf-8 -*-

import time
import logging
import json
import os
from typing import List, Dict, Any, Optional, Tuple

import pygame

class JoystickDetector:
    """
    游戏手柄检测工具，用于检测和识别已连接的游戏手柄设备
    """
    
    # 常见游戏手柄按键名称映射
    BUTTON_NAMES = {
        # Xbox 控制器
        "Xbox Controller": {
            0: "A",
            1: "B", 
            2: "X",
            3: "Y",
            4: "LB",
            5: "RB",
            6: "Back",
            7: "Start",
            8: "LS",  # 左摇杆按下
            9: "RS",  # 右摇杆按下
        },
        # Xbox 360 控制器
        "Xbox 360 Controller": {
            0: "A",
            1: "B", 
            2: "X",
            3: "Y",
            4: "LB",
            5: "RB",
            6: "Back",
            7: "Start",
            8: "LS",
            9: "RS",
        },
        # PS4 控制器
        "PS4 Controller": {
            0: "×",  # 叉
            1: "○",  # 圆
            2: "□",  # 方
            3: "△",  # 三角
            4: "L1",
            5: "R1", 
            6: "L2",
            7: "R2",
            8: "Share",
            9: "Options",
            10: "PS",  # PS按钮
            11: "L3",  # 左摇杆按下
            12: "R3",  # 右摇杆按下
        },
        # 通用蓝牙游戏手柄
        "Bluetooth Controller": {
            0: "A/×",
            1: "B/○",
            2: "X/□",
            3: "Y/△",
            4: "L1/LB",
            5: "R1/RB",
            6: "Back/Share",
            7: "Start/Options",
            8: "L3/LS",
            9: "R3/RS",
        },
        # 罗技控制器
        "Logitech Controller": {
            0: "A/×",
            1: "B/○",
            2: "X/□",
            3: "Y/△",
            4: "LB",
            5: "RB",
            6: "Back",
            7: "Start",
            8: "Home",
            9: "L3",
            10: "R3",
        },
        # 任天堂 Switch Pro 控制器
        "Nintendo Switch Pro Controller": {
            0: "B",
            1: "A", 
            2: "Y",
            3: "X",
            4: "L",
            5: "R", 
            6: "ZL",
            7: "ZR",
            8: "-",
            9: "+",
            10: "Home",
            11: "Capture",
            12: "LS",  # 左摇杆按下
            13: "RS",  # 右摇杆按下
        },
    }
    
    # 通用控制器映射（如果没有匹配到具体型号）
    GENERIC_BUTTON_NAMES = {
        0: "A/×",
        1: "B/○",
        2: "X/□",
        3: "Y/△",
        4: "LB/L1",
        5: "RB/R1",
        6: "Back/Select",
        7: "Start/Options",
        8: "Home/Guide",
        9: "LS/L3",
        10: "RS/R3",
    }
    
    def __init__(self) -> None:
        """
        初始化游戏手柄检测器
        """
        # 获取日志记录器而不是创建新的
        self.logger = logging.getLogger("JoystickDetector")
        
        # 初始化pygame
        pygame.init()
        pygame.joystick.init()
    
    def get_joystick_count(self) -> int:
        """
        获取已连接的游戏手柄数量
        
        :return: 游戏手柄数量
        """
        return pygame.joystick.get_count()
    
    def get_joystick_info(self, joystick_id: int = 0) -> Optional[Dict[str, Any]]:
        """
        获取指定游戏手柄的详细信息
        
        :param joystick_id: 游戏手柄ID，默认为0
        :return: 包含游戏手柄信息的字典，如果手柄不存在则返回None
        """
        if joystick_id >= self.get_joystick_count():
            self.logger.error(f"游戏手柄ID {joystick_id} 超出范围")
            return None
        
        try:
            joystick = pygame.joystick.Joystick(joystick_id)
            joystick.init()
            
            info = {
                "id": joystick_id,
                "name": joystick.get_name(),
                "guid": joystick.get_guid(),
                "num_axes": joystick.get_numaxes(),
                "num_buttons": joystick.get_numbuttons(),
                "num_hats": joystick.get_numhats(),
                "num_balls": joystick.get_numballs()
            }
            
            joystick.quit()
            return info
        except Exception as e:
            self.logger.error(f"获取游戏手柄 {joystick_id} 信息失败: {str(e)}")
            return None
    
    def list_joysticks(self) -> List[Dict[str, Any]]:
        """
        列出所有已连接的游戏手柄信息
        
        :return: 包含所有游戏手柄信息的列表
        """
        joysticks = []
        count = self.get_joystick_count()
        
        for i in range(count):
            info = self.get_joystick_info(i)
            if info:
                joysticks.append(info)
        
        return joysticks
    
    def wait_for_joystick(self, timeout: float = 0) -> Optional[Dict[str, Any]]:
        """
        等待游戏手柄连接
        
        :param timeout: 超时时间（秒），0表示无限等待
        :return: 首个检测到的游戏手柄信息，超时则返回None
        """
        start_time = time.time()
        
        while True:
            if timeout > 0 and time.time() - start_time > timeout:
                self.logger.warning(f"等待游戏手柄连接超时")
                return None
            
            # 强制刷新pygame事件队列
            pygame.event.pump()
            pygame.joystick.quit()
            pygame.joystick.init()
            
            count = self.get_joystick_count()
            if count > 0:
                self.logger.info(f"检测到 {count} 个游戏手柄")
                return self.get_joystick_info(0)
            
            time.sleep(0.5)
    
    def get_button_name(self, controller_name: str, button_id: int) -> str:
        """
        根据控制器名称和按钮ID获取按钮名称
        
        :param controller_name: 控制器名称
        :param button_id: 按钮ID
        :return: 按钮名称
        """
        # 尝试查找精确匹配的控制器名称
        for known_name, button_map in self.BUTTON_NAMES.items():
            if known_name.lower() in controller_name.lower():
                if button_id in button_map:
                    return button_map[button_id]
        
        # 如果没有精确匹配，使用通用名称
        if button_id in self.GENERIC_BUTTON_NAMES:
            return self.GENERIC_BUTTON_NAMES[button_id]
        
        # 最后使用ID作为名称
        return f"Button {button_id}"
    
    def load_custom_mapping(self, mapping_file: str) -> Dict[int, str]:
        """
        从文件加载自定义按钮映射
        
        :param mapping_file: 映射文件路径
        :return: 按钮ID到名称的映射字典
        """
        if not os.path.exists(mapping_file):
            self.logger.warning(f"映射文件 {mapping_file} 不存在")
            return {}
        
        try:
            with open(mapping_file, 'r', encoding='utf-8') as f:
                mapping = json.load(f)
            
            # 将字符串键转换为整数键
            int_mapping = {}
            for key, value in mapping.items():
                try:
                    int_mapping[int(key)] = value
                except ValueError:
                    self.logger.warning(f"忽略无效的按钮ID: {key}")
            
            self.logger.info(f"从 {mapping_file} 加载了 {len(int_mapping)} 个按钮映射")
            return int_mapping
        except Exception as e:
            self.logger.error(f"加载映射文件失败: {str(e)}")
            return {}
    
    def test_buttons(self, joystick_id: int = 0, timeout: float = 30, mapping_file: Optional[str] = None) -> Dict[int, str]:
        """
        测试游戏手柄按钮，让用户按下每个按钮并记录按钮名称
        
        :param joystick_id: 游戏手柄ID，默认为0
        :param timeout: 测试超时时间（秒）
        :param mapping_file: 自定义按钮映射文件路径，如果提供则优先使用
        :return: 按钮ID到按钮名称的映射字典
        """
        if joystick_id >= self.get_joystick_count():
            self.logger.error(f"游戏手柄ID {joystick_id} 超出范围")
            return {}
        
        try:
            joystick = pygame.joystick.Joystick(joystick_id)
            joystick.init()
            
            controller_name = joystick.get_name()
            button_count = joystick.get_numbuttons()
            button_map = {}
            
            # 加载自定义映射
            custom_mapping = {}
            if mapping_file:
                custom_mapping = self.load_custom_mapping(mapping_file)
            
            self.logger.info(f"开始测试游戏手柄 '{controller_name}' 的按钮")
            self.logger.info(f"请依次按下游戏手柄上的每个按钮，{timeout}秒后自动结束")
            
            start_time = time.time()
            running = True
            
            while running and time.time() - start_time < timeout:
                for event in pygame.event.get():
                    if event.type == pygame.JOYBUTTONDOWN and event.joy == joystick_id:
                        button_id = event.button
                        if button_id not in button_map:
                            # 获取按钮名称的优先级：
                            # 1. 自定义映射
                            # 2. 预定义的控制器特定映射
                            # 3. 通用映射
                            if button_id in custom_mapping:
                                button_name = custom_mapping[button_id]
                            else:
                                button_name = self.get_button_name(controller_name, button_id)
                            
                            button_map[button_id] = button_name
                            self.logger.info(f"检测到按钮: {button_id} - {button_name}")
                
                # 如果已经检测到所有按钮，提前结束
                if len(button_map) == button_count:
                    self.logger.info("所有按钮已测试完成")
                    running = False
                
                time.sleep(0.05)
            
            if len(button_map) < button_count:
                self.logger.warning(f"未能检测到所有按钮 ({len(button_map)}/{button_count})")
            
            joystick.quit()
            return button_map
        except Exception as e:
            self.logger.error(f"测试游戏手柄按钮失败: {str(e)}")
            return {}
    
    def create_button_mapping_file(self, button_map: Dict[int, str], output_file: str) -> bool:
        """
        创建按钮映射文件
        
        :param button_map: 按钮ID到名称的映射字典
        :param output_file: 输出文件路径
        :return: 是否成功创建文件
        """
        try:
            # 确保目录存在
            os.makedirs(os.path.dirname(os.path.abspath(output_file)), exist_ok=True)
            
            # 将整数键转换为字符串键
            str_mapping = {str(k): v for k, v in button_map.items()}
            
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(str_mapping, f, ensure_ascii=False, indent=2)
            
            self.logger.info(f"按钮映射已保存到: {output_file}")
            return True
        except Exception as e:
            self.logger.error(f"创建按钮映射文件失败: {str(e)}")
            return False
    
    def cleanup(self) -> None:
        """
        清理资源
        """
        pygame.joystick.quit()
        pygame.quit()
        self.logger.info("已清理游戏手柄检测器资源") 